import os
import requests
import time
from tqdm import tqdm
from flask import current_app
import gdown

# 模型下载链接
MODEL_URLS = {
    'u2net': 'https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    # U2Net的另一个较小的模型，如果完整模型下载太慢可以考虑使用这个
    'u2netp': 'https://drive.google.com/uc?id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW',
    # UNet模型链接可以根据需要添加
}

def download_file(url, filename):
    """
    下载文件并显示进度条
    
    Args:
        url (str): 下载链接
        filename (str): 保存的文件名
    """
    # 如果是Google Drive链接，使用gdown下载
    if "drive.google.com" in url:
        try:
            print(f"正在使用gdown从Google Drive下载模型: {filename}")
            gdown.download(url, filename, quiet=False)
            return True
        except Exception as e:
            print(f"使用gdown下载失败: {str(e)}，尝试使用requests下载")
            
    # 常规下载方式
    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))
    block_size = 1024  # 1 KB
    
    print(f"开始下载模型: {filename}")
    print(f"文件大小: {total_size/1024/1024:.2f} MB")
    
    progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True)
    
    with open(filename, 'wb') as file:
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    
    progress_bar.close()
    
    if total_size != 0 and progress_bar.n != total_size:
        print("下载出错，请重试!")
        return False
    
    print(f"模型下载完成: {filename}")
    return True

def generate_dummy_model(model_path, model_name):
    """
    生成一个只有几KB的虚拟模型文件，仅用于测试
    
    Args:
        model_path (str): 模型保存路径
        model_name (str): 模型名称
        
    Returns:
        bool: 是否成功生成
    """
    try:
        import torch
        import torch.nn as nn
        
        # 创建一个小型模型
        class DummyModel(nn.Module):
            def __init__(self):
                super(DummyModel, self).__init__()
                self.conv1 = nn.Conv2d(3, 10, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(10, 1, kernel_size=3, padding=1)
                
            def forward(self, x):
                x = torch.relu(self.conv1(x))
                x = self.conv2(x)
                return x
        
        model = DummyModel()
        
        # 保存模型
        torch.save(model.state_dict(), model_path)
        print(f"已生成测试用虚拟模型: {model_path}")
        return True
    except Exception as e:
        print(f"生成虚拟模型失败: {str(e)}")
        return False

def download_model_if_needed(model_name):
    """
    检查模型是否存在，如果不存在则下载
    
    Args:
        model_name (str): 模型名称 ('u2net', 'u2netp', 'unet')
        
    Returns:
        bool: 下载是否成功
    """
    # 确定模型文件路径
    if model_name == 'u2net':
        model_filename = 'u2net_model.pth'
    elif model_name == 'u2netp':
        model_filename = 'u2netp_model.pth'
    elif model_name == 'unet':
        model_filename = 'unet_model.pth'
    else:
        raise ValueError(f"不支持的模型名称: {model_name}")
    
    # 获取模型文件夹路径
    model_dir = os.path.join(os.getcwd(), 'models', 'weights')
    # 确保目录存在
    os.makedirs(model_dir, exist_ok=True)
    
    # 完整的模型文件路径
    model_path = os.path.join(model_dir, model_filename)
    
    # 检查模型是否已经存在
    if os.path.exists(model_path):
        print(f"模型已存在: {model_path}")
        return True
    
    # 创建一个虚拟模型用于测试（快速）
    use_dummy_model = os.environ.get('USE_DUMMY_MODEL', 'false').lower() == 'true'
    if use_dummy_model:
        print(f"使用虚拟模型代替下载 (环境变量USE_DUMMY_MODEL=true)")
        return generate_dummy_model(model_path, model_name)
    
    # 如果模型不存在，下载
    if model_name in MODEL_URLS:
        url = MODEL_URLS[model_name]
        print(f"模型不存在，准备下载: {model_name}")
        return download_file(url, model_path)
    else:
        print(f"没有找到模型的下载链接: {model_name}")
        return False 